package resolvers

import (
	"testing"

	"github.com/stackrox/rox/generated/storage"
	"github.com/stackrox/rox/pkg/set"
	"github.com/stretchr/testify/assert"
)

func mapImagesToVulnerabilityCounter(images []*storage.Image) *VulnerabilityCounterResolver {
	var vulns []*storage.EmbeddedVulnerability
	seenVulns := set.NewStringSet()
	for _, image := range images {
		for _, component := range image.GetScan().GetComponents() {
			for _, vuln := range component.GetVulns() {
				if seenVulns.Contains(vuln.GetCve()) {
					continue
				}
				seenVulns.Add(vuln.GetCve())
				vulns = append(vulns, vuln)
			}
		}
	}
	return mapVulnsToVulnerabilityCounter(vulns)
}

func TestMapImagesToVulnerabilityCounterResolver(t *testing.T) {
	images := testImages()
	images[0].Scan.Components = append(images[0].Scan.Components, &storage.EmbeddedImageScanComponent{
		Name:    "comp1",
		Version: "0.9",
		Vulns: []*storage.EmbeddedVulnerability{
			{
				Cve: "cve-2018-4",
				SetFixedBy: &storage.EmbeddedVulnerability_FixedBy{
					FixedBy: "1.1",
				},
				Cvss:     10.0,
				Severity: storage.VulnerabilitySeverity_CRITICAL_VULNERABILITY_SEVERITY,
			},
		},
	})
	images[1].Scan.Components = append(images[1].Scan.Components, &storage.EmbeddedImageScanComponent{
		Name:    "comp1",
		Version: "0.9",
		Vulns: []*storage.EmbeddedVulnerability{
			{
				Cve: "cve-2018-6",
				SetFixedBy: &storage.EmbeddedVulnerability_FixedBy{
					FixedBy: "1.1",
				},
				Cvss:     1.9,
				Severity: storage.VulnerabilitySeverity_LOW_VULNERABILITY_SEVERITY,
			},
		},
	})

	counterResolver := mapImagesToVulnerabilityCounter(images)
	assert.Equal(t, int32(7), counterResolver.all.total)
	assert.Equal(t, int32(3), counterResolver.all.fixable)
	assert.Equal(t, int32(2), counterResolver.low.total)
	assert.Equal(t, int32(1), counterResolver.low.fixable)
	assert.Equal(t, int32(1), counterResolver.moderate.total)
	assert.Equal(t, int32(0), counterResolver.moderate.fixable)
	assert.Equal(t, int32(2), counterResolver.important.total)
	assert.Equal(t, int32(0), counterResolver.important.fixable)
	assert.Equal(t, int32(2), counterResolver.critical.total)
	assert.Equal(t, int32(2), counterResolver.critical.fixable)
}
